import cv2
import numpy as np
import os
from matplotlib import pyplot as plt
# 预设RGB值，在后面处理时筛选出蓝色区域
lower_blue = np.array([160, 80, 19])
upper_blue = np.array([250, 150, 250])
lower_yellow = np.array([15, 55, 55])
upper_yellow = np.array([50, 255, 255])
# 在调用cv2库的时候偶尔会报错，经查询是路径的问题，我这里图片全部使用绝对路径


def get_path(dirname):  # 图片所在目录名
    path = os.path.dirname(__file__)  # 获取当前路径
    img_names = os.listdir(os.path.join(path+'/'+dirname+'/'))  # 目录下所有文件名
    img_path = os.path.join(path + '/' + dirname + '/')  # 目录下所有文件绝对路径
    return img_names, img_path  # 返回文件名和绝对路径


def enhance_contrast(image):  # 图像增强
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    img_tophat = cv2.morphologyEx(image, cv2.MORPH_TOPHAT, kernel)
    img_blackhat = cv2.morphologyEx(
        image, cv2.MORPH_BLACKHAT, kernel)  # 这里主要将图片腐蚀话，过滤掉较小的噪声
    image_plus_tophat = cv2.add(image, img_tophat)
    image_plus_blackhat_minus_blackhat = cv2.subtract(
        image_plus_tophat, img_blackhat)
    return image_plus_blackhat_minus_blackhat


def img_processing(img, ret_path, filename):  # 进一步图像处理，写的比较乱，图像处理步骤可以重新调整下
    img = cv2.imread(img)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15))
    img_ = enhance_contrast(img)
    cv2.imwrite(ret_path+'001.jpg', img_)
    hsv = cv2.cvtColor(img_, cv2.COLOR_BGR2HSV)
    mask_blue = cv2.inRange(img, lower_blue, upper_blue)
    cv2.imwrite(ret_path+'002.jpg', mask_blue)
    mask_yellow = cv2.inRange(hsv, lower_yellow, upper_yellow)
    cv2.imwrite(ret_path+'003.jpg', mask_yellow)
    output = cv2.bitwise_and(hsv, hsv, mask=mask_blue)
    cv2.imwrite(ret_path+'004.jpg', output)
    # 根据阈值找到对应颜色
    closed = cv2.morphologyEx(output, cv2.MORPH_CLOSE, kernel)
    cv2.imwrite(ret_path+'005.jpg', closed)
    closed = cv2.morphologyEx(closed, cv2.MORPH_CLOSE, kernel)
    cv2.imwrite(ret_path+'006.jpg', closed)
    closed_ = cv2.dilate(closed, None, iterations=4)  # 这里用一个较大的值又对图片腐蚀了一下
    cv2.imwrite(ret_path+'007.jpg', closed_)
    gray = cv2.cvtColor(closed_, cv2.COLOR_BGR2GRAY)
    cv2.imwrite(ret_path+'008.jpg', gray)
    ret, binary = cv2.threshold(gray, 100, 255, cv2.THRESH_BINARY)
    cv2.imwrite(ret_path+'009.jpg', ret)
    cv2.imwrite(ret_path+'010.jpg', binary)

    ctrs, hier = cv2.findContours(
        binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  # 对处理好的图片查找边界
    box_filter = []
    a = sorted(ctrs, key=cv2.contourArea, reverse=True)[0]
    for i in range(len(ctrs)):
        c = sorted(ctrs, key=cv2.contourArea, reverse=True)[i]
        # compute the rotated bounding box of the largest contour
        rect = cv2.minAreaRect(c)  # 最小外界矩形，到这步很多不规则噪声就会过滤掉了
        box = np.int0(cv2.boxPoints(rect))
        Xs = [i[0] for i in box]
        Ys = [i[1] for i in box]
        x1 = min(Xs)
        x2 = max(Xs)
        y1 = min(Ys)
        y2 = max(Ys)
        hight = y2 - y1
        width = x2 - x1
        # 很多时候会有一些其他矩形干扰，这里把长宽比3/1的找出来，并且面积大于800
        if width/hight < 4 and width/hight > 2 and cv2.contourArea(box) > 800:
            box_filter.append(box)  # 多个车牌的时候把所有符合要求的放入到list中
            cv2.drawContours(img, [box], -1, (0, 0, 255), 1)  # 绘制出矩形框，
            cropImg = img[y1:y1 + hight, x1:x1 + width]  # 定位输出区域
            cv2.imwrite(ret_path+'plate_num/'+filename+'.jpg',
                        cropImg)  # 将输出保持为图片存在文件夹下
        else:
            pass
    if box_filter == []:  # 对于有变形的无法识别出车牌，这里将最大矩形面积存为图片
        rect0 = cv2.minAreaRect(a)
        box0 = np.int0(cv2.boxPoints(rect0))
        cv2.drawContours(img, [box0], -1, (0, 0, 255), 1)
        Xs = [i[0] for i in box0]
        Ys = [i[1] for i in box0]
        x1 = min(Xs)
        x2 = max(Xs)
        y1 = min(Ys)
        y2 = max(Ys)
        hight = y2 - y1
        width = x2 - x1
        cropImg = img[y1:y1 + hight, x1:x1 + width]
        cv2.imwrite(ret_path+'plate_num/' + filename+'.jpg', cropImg)
    else:
        pass
    return img
def split_img():
    img=cv2.imread('plate_num/timg5.jpeg.jpg')#读取上一步定位的车牌信息
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img_thre = gray
    ret, thresh1=cv2.threshold(img_thre, 130, 255, cv2.THRESH_BINARY)
    #图片二值化处理，原先多个通道的现在只有一个 0-黑色 255-白色 
    print(thresh1.shape)
    # 以下为识别车牌中的字符
    w_point=[]
    w_point0=[]
    height = thresh1.shape[0] #获取车牌高
    width = thresh1.shape[1] #获取车牌长
    for i in range(height):#去掉w_point数值小于X的行
        a = 0
        for j in range(width):
            if thresh1[i][j] == 255: 
                a+=1
        if 300>a>50 :#如果一行中白色元素在300/5之间则保留，其余行全部转为黑色
            continue
        else:
            for b in range(width):
                thresh1[i][b]=0
    for i in range(width):#去掉w_point数值小于X的列
        a = 0
        for j in range(height):
            if thresh1[j][i] == 255:
                a+=1
        if 300>a>=20 :#如果一列中白色元素在300/5之间则保留，其余列全部转为黑色
            continue
        else:
            for b in range(height):
                thresh1[b][i] = 0
     #因为是按列分割，所以行不变，对列进行分析，如果一列上白元素为0那肯定是分割点           
    for i in range(width):#对于刚才处理完的图片，重新查询每一列白色
        a = 0
        for j in range(height):
            if thresh1[j][i] == 255:
                a+=1
        w_point.append(a)#生成一两list用来存储每一列白元素的数量
        letter = []
        n = 0
        for i in range(len(w_point) - 1):
            if w_point[i] == 0 and w_point[i + 1] != 0:
                letter.append(i)#生成一个list，保存了从哪列开始有白元素
    #针对demo letter列表为[6, 23, 80, 164, 225, 282, 344, 399, 453]，也就是从6，23，80....开始有白元素
    print(letter)
    for i in range(len(letter)-1):#直接安装letter点对图像进行分割
        start=letter[i]
        end=letter[i+1]
        i=i+1
        print(start,end)
        if end-start>10# 去掉边框及误差
        	img_=thresh1[:,start:end]
        # cv2.imshow('img',img_)
        	cv2.imwrite('spilt_num/'+'%s.jpg'%i,img_)

            
if __name__ == "__main__":
    num = 0
    path = os.path.dirname(__file__)
    img_names = os.listdir(os.path.join(path+'/img/'))  # 目录下所有文件名
    img_path = os.path.join(path + '/img/')  # 目录下所有文件绝对路径
    ret_path = os.path.join(path + '/img_ret/')
    img = img_processing(img_path+img_names[0], ret_path, "final_ret")
